import torch
import numpy as np
from scipy.interpolate import CubicSpline
from matplotlib.animation import FuncAnimation
import random

class PeopleController:
    def __init__(self, num_points = 10):
        # 生成随机控制点（使用PyTorch生成张量）
        self.num_points = num_points
        # 在参数化时间轴上的移动比例
        self.speed = 0.0005 * random.uniform(0.1, 1.0)  # 运动速度
        self.direction = 1   # 运动方向
        self.current_t = 0.0
        self.trajectory = self.gen_trajectory()

    def gen_trajectory(self):
        control_points = torch.rand(self.num_points, 2) * 10 - 5  # 生成-5到5之间的二维坐标
        # 转换为numpy进行插值（保留PyTorch梯度计算能力）
        np_points = control_points.detach().numpy()
        t_param = np.linspace(0, 1, self.num_points)
        # 创建三次样条插值函数
        cs_x = CubicSpline(t_param, np_points[:, 0], bc_type='natural')
        cs_y = CubicSpline(t_param, np_points[:, 1], bc_type='natural')
        # 生成密集轨迹点
        t_dense = np.linspace(0, 1, 1000)
        trajectory = torch.tensor(np.column_stack((cs_x(t_dense), cs_y(t_dense))), dtype=torch.float32)
        return trajectory
        
    def update(self):
        self.current_t += self.speed * self.direction
        # 到达边界时反转方向
        if self.current_t >= 1.0 or self.current_t <= 0.0:
            self.direction *= -1
            self.current_t = torch.clamp(torch.tensor(self.current_t), 0.0, 1.0).item()
        # 计算当前位置
        idx = int(self.current_t * (len(self.trajectory)-1))
        return self.trajectory[idx]
    
class BatchPeopleController:
    def __init__(self, num_objects=128, num_points=10, device='cpu'):
        self.num_objects = num_objects
        self.num_points = num_points
        self.device = device
        
        # 批量生成初始参数
        self.speed = torch.rand(num_objects, device=device) * 0.0005 + 0.0005  # [0.0005, 0.001]
        self.direction = torch.ones(num_objects, device=device)
        self.current_t = torch.zeros(num_objects, device=device)
        
        # 生成轨迹和切线方向
        self.trajectories, self.tangents = self.gen_trajectories()
    
    def gen_trajectories(self):
        """批量生成轨迹和切线方向"""
        control_points = torch.rand(self.num_objects, self.num_points, 2, device=self.device) * 10 - 5
        
        traj_list, tangent_list = [], []
        for i in range(self.num_objects):
            # 生成位置轨迹
            np_points = control_points[i].cpu().numpy()
            t_param = np.linspace(0, 1, self.num_points)
            
            # 三次样条及其导数
            cs_x = CubicSpline(t_param, np_points[:, 0], bc_type='natural')
            cs_y = CubicSpline(t_param, np_points[:, 1], bc_type='natural')
            
            # 生成密集采样点
            t_dense = np.linspace(0, 1, 1000)
            traj = torch.tensor(np.column_stack((cs_x(t_dense), cs_y(t_dense))), 
                               dtype=torch.float32, device=self.device)
            
            # 计算切线方向（参数空间的导数）
            dx = cs_x(t_dense, 1)
            dy = cs_y(t_dense, 1)
            tangent = torch.tensor(np.column_stack((dx, dy)), 
                                  dtype=torch.float32, device=self.device)
            
            traj_list.append(traj)
            tangent_list.append(tangent)
            
        return torch.stack(traj_list), torch.stack(tangent_list)

    def update(self):
        """返回位置和四元数朝向"""
        # 更新时间参数
        self.current_t += self.speed * self.direction
        
        # 边界处理
        overflow_mask = (self.current_t >= 1.0) | (self.current_t <= 0.0)
        self.direction[overflow_mask] *= -1
        self.current_t = torch.clamp(self.current_t, 0.0, 1.0)
        
        # 计算采样索引
        indices = (self.current_t * (self.trajectories.shape[1] - 1)).long()
        
        # 获取当前位置 [N, 2]
        positions = self.trajectories[torch.arange(self.num_objects, device=self.device), indices]
        
        # 计算朝向四元数 ----------------------------------------------
        # 获取原始切线方向 [N, 2]
        tangents = self.tangents[torch.arange(self.num_objects, device=self.device), indices]
        
        # 根据运动方向调整切线方向
        adjusted_tangents = tangents * self.direction.unsqueeze(-1)
        
        # 处理零向量问题
        lengths = torch.norm(adjusted_tangents, dim=1, keepdim=True)
        mask = (lengths < 1e-6).squeeze()
        adjusted_tangents[mask] = torch.tensor([1.0, 0.0], device=self.device)  # 默认向前
        adjusted_tangents = adjusted_tangents / (lengths + 1e-6)  # 归一化
        
        # 计算旋转角度（绕Z轴）
        theta = torch.atan2(adjusted_tangents[:, 1], adjusted_tangents[:, 0])
        
        # 转换为四元数 [N, 4]
        half_theta = theta / 2
        w = torch.cos(half_theta)
        z = torch.sin(half_theta)
        quats = torch.zeros(self.num_objects, 4, device=self.device)
        quats[:, 0] = w  # 四元数w分量
        quats[:, 3] = z  # 四元数z分量（绕Z轴旋转）
        
        return positions, quats